import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
os.environ['TORCH_MODEL_ZOO'] = 'https://mirrors.tuna.tsinghua.edu.cn/pytorch-wheels/torch_vision/'
from sklearn.metrics import accuracy_score
import time
from datetime import timedelta
import matplotlib.pyplot as plt
import sys
import matplotlib
from matplotlib.font_manager import FontProperties
import pywt
import cv2
from torchvision import models

try:
    font = FontProperties(fname=r"C:\Windows\Fonts\msyh.ttc")
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei']
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
except:
    plt.rcParams['font.sans-serif'] = ['SimHei', 'KaiTi', 'Arial Unicode MS']
    plt.rcParams['axes.unicode_minus'] = False

matplotlib.use('TkAgg', force=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义CWT数据集类
class CWTDataset(Dataset):
    def __init__(self, current_root, voltage_root, transform=None, target_size=(224, 224), scales=None):
        """
        Args:
            current_root (string): 电流数据的根目录路径（包含class0到class8文件夹）
            voltage_root (string): 电压数据的根目录路径（包含class0到class8文件夹）
            transform (callable, optional): 可选的图像变换/增强函数
            target_size (tuple): 目标图像尺寸 (height, width)
            scales (list): 小波变换的尺度参数
        """
        self.current_root = current_root
        self.voltage_root = voltage_root
        self.transform = transform
        self.target_size = target_size

        if scales is None:
            self.scales = np.arange(1, 64)

        self.current_file_paths = []
        self.voltage_file_paths = []
        self.labels = []

        classes = sorted([d for d in os.listdir(current_root) if os.path.isdir(os.path.join(current_root, d))])
        self.class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
        self.num_classes = len(classes)

        for cls_name in classes:
            current_cls_dir = os.path.join(current_root, cls_name)
            voltage_cls_dir = os.path.join(voltage_root, cls_name)

            current_files = sorted([f for f in os.listdir(current_cls_dir) if f.endswith('.npy')])
            voltage_files = sorted([f for f in os.listdir(voltage_cls_dir) if f.endswith('.npy')])

            common_files = set(current_files) & set(voltage_files)

            for file_name in common_files:
                current_file_path = os.path.join(current_cls_dir, file_name)
                voltage_file_path = os.path.join(voltage_cls_dir, file_name)

                self.current_file_paths.append(current_file_path)
                self.voltage_file_paths.append(voltage_file_path)
                self.labels.append(self.class_to_idx[cls_name])

    def __len__(self):
        return len(self.current_file_paths)

    def __getitem__(self, idx):
        current_data = np.load(self.current_file_paths[idx])
        voltage_data = np.load(self.voltage_file_paths[idx])

        # 降采样 - 每隔downsample_factor个点取一个点  慎改，非常影响原信号频率特性！！！
        downsample_factor = 1
        current_data = current_data[::downsample_factor]
        voltage_data = voltage_data[::downsample_factor]

        cwt_current, _ = pywt.cwt(current_data, self.scales, 'morl')
        cwt_voltage, _ = pywt.cwt(voltage_data, self.scales, 'morl')

        # 取绝对值并转换为dB尺度，增强对比度
        cwt_current = np.abs(cwt_current)
        cwt_voltage = np.abs(cwt_voltage)

        # resize
        cwt_current = cv2.resize(cwt_current, self.target_size[::-1], interpolation=cv2.INTER_AREA)
        cwt_voltage = cv2.resize(cwt_voltage, self.target_size[::-1], interpolation=cv2.INTER_AREA)

        # 将两个时频图堆叠为2通道图像
        image = np.stack([cwt_current, cwt_voltage], axis=0)  # 形状: (2, H, W)

        image = torch.from_numpy(image).float()

        # 简单归一化到[0,1]
        min_val = image.amin(dim=(1, 2), keepdim=True)
        max_val = image.amax(dim=(1, 2), keepdim=True)
        image = (image - min_val) / (max_val - min_val + 1e-8)

        if self.transform:
            image = self.transform(image)

        label = self.labels[idx]
        return image, label

# 定义ResNet模型（修改为2通道输入）
def get_resnet_model(num_classes=9, model_name='resnet18'):
    if model_name == 'resnet18':
        model = models.resnet18(pretrained=True)
    elif model_name == 'resnet34':
        model = models.resnet34(pretrained=True)
    elif model_name == 'resnet50':
        model = models.resnet50(pretrained=True)
    else:
        raise ValueError(f"Unsupported model: {model_name}")

    # 修改第一层卷积，输入通道从3改为2
    original_first_conv = model.conv1
    model.conv1 = nn.Conv2d(2, original_first_conv.out_channels,
                            kernel_size=original_first_conv.kernel_size,
                            stride=original_first_conv.stride,
                            padding=original_first_conv.padding,
                            bias=original_first_conv.bias)

    # 修改最后的全连接层，输出类别数为num_classes
    num_ftrs = model.fc.in_features
    model.fc = nn.Linear(num_ftrs, num_classes)

    return model

# 计算准确率
def compute_accuracy(model, data_loader):
    model.eval()
    all_labels = []
    all_preds = []
    with torch.no_grad():
        for images, labels in data_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            all_labels.extend(labels.cpu().numpy())
            all_preds.extend(predicted.cpu().numpy())
    return accuracy_score(all_labels, all_preds)

# 训练函数
def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs):
    plt.ion()
    fig, axs = plt.subplots(1, 3, figsize=(18, 5))

    train_losses = []
    test_losses = []
    train_accuracies = []
    test_accuracies = []

    best_accuracy = 0.0
    start_time = time.time()
    save_dir = r".\models"
    os.makedirs(save_dir, exist_ok=True)
    best_model_path = os.path.join(save_dir, "best_model_resnet_cwt.pt")

    fig.suptitle('训练过程实时监控', fontsize=16)

    line_train_loss, = axs[0].plot([], [], 'b-', label='训练损失')
    line_test_loss, = axs[0].plot([], [], 'y-', label='测试损失')
    line_train_acc, = axs[1].plot([], [], 'g-o', label='训练准确率')
    line_test_acc, = axs[2].plot([], [], 'r-o', label='测试准确率')

    for ax in axs:
        ax.grid(True)
        ax.set_xlabel('Epoch')
        ax.legend(loc='best')  # 添加图例

    axs[0].set_ylabel('Loss')
    axs[0].set_title('训练和测试损失')
    axs[1].set_ylabel('Accuracy')
    axs[1].set_title('训练集准确率')
    axs[1].set_ylim(0, 1.0)
    axs[2].set_ylabel('Accuracy')
    axs[2].set_title('测试集准确率')
    axs[2].set_ylim(0, 1.0)

    plt.tight_layout()
    plt.subplots_adjust(top=0.85)  # 为标题留出空间

    for epoch in range(num_epochs):
        model.train()
        epoch_train_loss = 0.0

        all_train_labels = []
        all_train_preds = []

        for i, (images, labels) in enumerate(train_loader):
            images = images.to(device)
            labels = labels.to(device)

            # 前向传播
            outputs = model(images)
            loss = criterion(outputs, labels)

            _, predicted = torch.max(outputs.data, 1)
            all_train_labels.extend(labels.cpu().numpy())
            all_train_preds.extend(predicted.cpu().numpy())

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            epoch_train_loss += loss.item()

        avg_train_loss = epoch_train_loss / len(train_loader)
        train_losses.append(avg_train_loss)

        train_accuracy = accuracy_score(all_train_labels, all_train_preds)
        train_accuracies.append(train_accuracy)

        model.eval()
        epoch_test_loss = 0.0
        all_test_labels = []
        all_test_preds = []
        with torch.no_grad():
            for images, labels in test_loader:
                images = images.to(device)
                labels = labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)
                epoch_test_loss += loss.item()

                _, predicted = torch.max(outputs.data, 1)
                all_test_labels.extend(labels.cpu().numpy())
                all_test_preds.extend(predicted.cpu().numpy())

        avg_test_loss = epoch_test_loss / len(test_loader)
        test_losses.append(avg_test_loss)
        test_accuracy = accuracy_score(all_test_labels, all_test_preds)
        test_accuracies.append(test_accuracy)

        print(f'Epoch [{epoch + 1}/{num_epochs}], 训练损失: {avg_train_loss:.4f}, 测试损失: {avg_test_loss:.4f}, '
              f'训练准确率: {train_accuracy:.4f}, 测试准确率: {test_accuracy:.4f}')

        if test_accuracy > best_accuracy:
            best_accuracy = test_accuracy
            model.eval()
            # 创建一个示例输入（2通道，224x224图像）
            example_input = torch.randn(1, 2, 224, 224).to(device)
            traced_model = torch.jit.trace(model, example_input)
            traced_model.save(best_model_path)
            print(f"发现新的最佳模型! 测试准确率: {test_accuracy:.4f}, 已保存至: {best_model_path}")

        epochs_range = range(1, len(train_losses) + 1)

        line_train_loss.set_data(epochs_range, train_losses)
        line_test_loss.set_data(epochs_range, test_losses)
        axs[0].relim()
        axs[0].autoscale_view()

        line_train_acc.set_data(epochs_range, train_accuracies)
        axs[1].relim()
        axs[1].autoscale_view()

        line_test_acc.set_data(epochs_range, test_accuracies)
        axs[2].relim()
        axs[2].autoscale_view()

        fig.canvas.draw()
        fig.canvas.flush_events()

        # 添加短暂延迟
        time.sleep(0.05)

    total_time = time.time() - start_time
    print(f'\n训练完成! 总时间: {timedelta(seconds=total_time)}')
    print(f'最高测试准确率: {best_accuracy:.4f}')

    plt.savefig(os.path.join(save_dir, 'training_metrics.png'))
    plt.ioff()
    plt.show()

    return best_accuracy

# 主程序
if __name__ == "__main__":
    # 设置路径
    current_train_root = r"D:\A_90762\atl\Current\train"
    voltage_train_root = r"D:\A_90762\atl\Voltage\train"
    current_test_root = r"D:\A_90762\atl\Current\test"
    voltage_test_root = r"D:\A_90762\atl\Voltage\test"

    # 创建数据集实例
    train_dataset = CWTDataset(
        current_root=current_train_root,
        voltage_root=voltage_train_root,
        target_size=(224, 224)
    )

    test_dataset = CWTDataset(
        current_root=current_test_root,
        voltage_root=voltage_test_root,
        target_size=(224, 224)
    )

    # 创建数据加载器
    batch_size = 16
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    num_classes = train_dataset.num_classes
    model = get_resnet_model(num_classes=num_classes, model_name='resnet18').to(device)

    learning_rate = 0.001
    num_epochs = 100

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    print(f"训练集大小: {len(train_dataset)}")
    print(f"测试集大小: {len(test_dataset)}")
    print(f"类别数量: {num_classes}")

    print("开始训练模型...")
    best_acc = train_model(
        model,
        train_loader,
        test_loader,
        criterion,
        optimizer,
        num_epochs
    )
    print("最佳模型已保存为TorchScript格式")

